from nnunet.network_architecture.generic_UNet import Generic_UNet
from nnunet.network_architecture.initialization import InitWeights_He
from torch import nn

# Reference from nnunet.training.network_training.nnUNetTrainerV2
norm_op_kwargs = {"eps": 1e-5, "affine": True}
dropout_op_kwargs = {"p": 0, "inplace": True}
net_nonlin_kwargs = {"negative_slope": 1e-2, "inplace": True}
class FullRes3D(Generic_UNet):
    def __init__(self) -> None:
        super(FullRes3D, self).__init__(
            1,
            32,
            2,
            5,
            2,
            2,
            nn.Conv3d,
            nn.InstanceNorm3d,
            norm_op_kwargs,
            nn.Dropout3d,
            dropout_op_kwargs,
            nn.LeakyReLU,
            net_nonlin_kwargs,
            True,
            False,
            lambda x: x,
            InitWeights_He(1e-2),
            [[1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]],
            [[1, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]],
            False,
            True,
            True,
        )
        # 做一些网络的初始化工作，优化预测性能
        # self.do_ds = False
